package edu.zju.gis.dbfg.model.tile;

import edu.zju.gis.dbfg.common.constant.Seprator;
import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.base.Pipeline;
import edu.zju.gis.dbfg.model.base.PipelineBuilder;
import edu.zju.gis.dbfg.model.base.VectorTileBuilder;
import edu.zju.gis.dbfg.model.io.geojson.GeoJsonTileBuilder;
import edu.zju.gis.dbfg.model.io.pbf.MVTTileBuilder;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import edu.zju.gis.dbfg.model.util.TileUtil;
import edu.zju.gis.dbfg.model.util.ZLevelInfo;
import lombok.Getter;
import lombok.Setter;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.*;
import org.geotools.geometry.jts.JTS;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.*;
import org.locationtech.jts.geom.Polygon;
import org.locationtech.jts.io.WKTReader;
import org.opengis.referencing.FactoryException;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.awt.*;
import java.io.Serializable;
import java.util.*;
import java.util.List;

/**
 * @author Hu
 * @date 2019/10/22
 **/
@Getter
@Setter
public class IteratorVectorTileClipper implements Serializable {

    private final static Logger logger = LoggerFactory.getLogger(VectorTileClipper.class);
    private final static int BUFFER = 1;
    private final static int SCREEN_TILE_SIZE = 256;
    private final static int OVERSAMPLE_FACTOR = 1;

    private VectorTileClipperArgs mArgs;
    private CoordinateReferenceSystem targetCRS;
    private CoordinateReferenceSystem sourceCRS;
    private PyramidConfig targetPyramidConfig;
    private MathTransform stTransform;
    private String datasource;


    private class TileBuilder implements PairFlatMapFunction<Tuple2<TileID, Iterable<GeoItem>>, TileID, GeoItem> {
        @Override
        public Iterator<Tuple2<TileID, GeoItem>> call(Tuple2<TileID, Iterable<GeoItem>> tileIDIterableTuple2) throws Exception {

            List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();

            TileID tileID = tileIDIterableTuple2._1;

            if (!tileIDIterableTuple2._2.iterator().hasNext()) {
                logger.warn(String.format("tile[%d, %d, %d] abort with total count [%d]", tileID.getzLevel(), tileID.getX(), tileID.getY(), 0));
                return result.iterator();
            }

            Iterator<GeoItem> geoItemIterator = tileIDIterableTuple2._2.iterator();
            Envelope tileEnvelope = TileUtil.createTileBox(tileID, targetPyramidConfig);
            Pipeline preProcessPipeline = getPipeline(CRS.parseWKT(mArgs.getTargetCRS()), BUFFER, tileEnvelope, true, false, false);

            VectorTileBuilder vectorTileBuilder;
            if (mArgs.getFileType().equals("geojson")) {
                vectorTileBuilder = new GeoJsonTileBuilder(mArgs.getOutDir(), tileID, targetPyramidConfig);
            } else {
                vectorTileBuilder = new MVTTileBuilder(mArgs.getOutDir(), tileID, tileEnvelope, mArgs.getLayerName());
            }

            TileID upperTileID = new TileID();
            upperTileID.setzLevel(tileID.getzLevel()-1);
            upperTileID.setX(tileID.getX()/2);
            upperTileID.setY(tileID.getY()/2);

            /**
             * 将之前被分割的图形union回来
             */
            Map<String, GeoItem> unionedGeoItem = new HashMap<>();

            while(geoItemIterator.hasNext()) {
                GeoItem geoItem = geoItemIterator.next();
                if (!unionedGeoItem.containsKey(geoItem.getFid())) {
                    unionedGeoItem.put(geoItem.getFid(), geoItem);
                } else {
                    Geometry g = onlyPolygon(geoItem.getGeometry());
                    Geometry og = onlyPolygon(unionedGeoItem.get(geoItem.getFid()).getGeometry());
                    if (og == null) continue;
                    Geometry ug = union(g, og, Math.max(tileEnvelope.getWidth()/(16*SCREEN_TILE_SIZE), tileEnvelope.getHeight()/(16*SCREEN_TILE_SIZE)), 0, 2);
                    if (ug == null) {
                        continue;
                    }
                    geoItem.setGeometry(ug);
                    // Map<String, Object> prop = geoItem.getProp();
                    // prop.put("k", (Double)prop.getOrDefault("k", 0d) + (Double)unionedGeoItem.get(geoItem.getFid()).getProp().getOrDefault("k", 0d));
                    // geoItem.setProp(prop);
                    unionedGeoItem.put(geoItem.getFid(), geoItem);
                }
            }


//            /**
//             * 构建R树
//             */
//             long startTime = System.currentTimeMillis();
//            Set<String> rKeys = unionedGeoItem.keySet();
//            for (String rKey: rKeys) {
//                GeoItem geoItem = unionedGeoItem.get(rKey);
//                if (geoItem.getGeometry().isValid()) {
//                    screenTile.insert(geoItem);
//                }
//            }
//             long endTime = System.currentTimeMillis();
//             logger.info("构建R树用时: " + (endTime - startTime) + ", 空间对象个数: " + screenTile.getGis().size());

//            /**
//             * 去除空间上重复的图斑
//             */
//            startTime = System.currentTimeMillis();
//            Iterator<String> mKeys = unionedGeoItem.keySet().iterator();
//            List<GeoItem> cleanItem = new ArrayList<>();
//            while (mKeys.hasNext()) {
//                String mkey = mKeys.next();
//                GeoItem geoItem = unionedGeoItem.get(mkey);
//                if (geoItem.getGeometry().isValid()) {
//                    screenTile.check(geoItem);
//                } else {
//                    cleanItem.add(geoItem);
//                }
//            }
//            cleanItem.addAll(screenTile.getSettedGeoItems());
//             endTime = System.currentTimeMillis();
//             logger.info("空间过滤用时: " + (endTime - startTime) + ", 剩余空间对象个数: " + cleanItem.size());
//             Iterator<String> keys = unionedGeoItem.keySet().iterator();

            List<Tuple2<TileID, GeoItem>> temp = new ArrayList<>();

            Set<String> rKeys = unionedGeoItem.keySet();
            for (String key: rKeys) {
//                String key = keys.next();
                GeoItem geoItem = unionedGeoItem.get(key);
                // Geometry clippedGeom = clipPipeline.execute(geoItem.getGeometry());
                Geometry finalGeom = preProcessPipeline.execute(geoItem.getGeometry());
                if (!finalGeom.isEmpty()) {
                    finalGeom = onlyPolygon(finalGeom);
                    if (finalGeom == null) continue;
                    vectorTileBuilder.addFeature(mArgs.getLayerName(), geoItem.getFid(), geoItem.getFid(),
                            finalGeom, new HashMap<>());
                    geoItem.setGeometry(finalGeom);
                    temp.add(new Tuple2<>(upperTileID, geoItem));
                }
            }

            vectorTileBuilder.build();
            logger.info(String.format("tile[%d, %d, %d] success with total count [%d]", tileID.getzLevel(), tileID.getX(), tileID.getY(), temp.size()));

            /**
             * 处理数据用于上一层计算，这时就不需要clip了，只需要进行抽稀即可
             */
            if (tileID.getzLevel() > mArgs.getZMin()) {
                Envelope upperLevelEnvelope = TileUtil.createTileBox(upperTileID, targetPyramidConfig);
                Pipeline upperLevelSimplifyPipeline = getPipeline(CRS.parseWKT(mArgs.getTargetCRS()), BUFFER, upperLevelEnvelope, false, false, true);
                for (int i=0; i<temp.size(); i++) {
                    GeoItem gi = temp.get(i)._2;
                    TileID ti = temp.get(i)._1;
                    Geometry g = upperLevelSimplifyPipeline.execute(gi.getGeometry());
                    if (g.isEmpty()) {
                        continue;
                    }
                    g = onlyPolygon(g);
                    if (g == null) {
                        continue;
                    }
                    gi.setGeometry(g);
                    result.add(new Tuple2<>(ti, gi));
                }
            } else {
                result.add(new Tuple2<>(null, null));
            }

            return result.iterator();
        }

    }

    public void init(String[] args) throws Exception {
        mArgs = VectorTileClipperArgs.initArgs(args);
        if (mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(VectorTileClipper.class, "main()", msg, args);
        }

        targetCRS = CRS.parseWKT(mArgs.getTargetCRS());
        sourceCRS = CRS.parseWKT(mArgs.getSourceCRS());
        targetPyramidConfig = new PyramidConfig.PyramidConfigBuilder()
                .setBaseMapEnv(-20026376.39, 20026376.39, -20048966.10, 20048966.10)
                .setCrs(targetCRS)
                .setzLevelRange(mArgs.getZMin(), mArgs.getZMax())
                .build();

        stTransform = CRS.findMathTransform(sourceCRS, targetCRS);
        datasource = mArgs.getInput();
    }

    /**
     * id => pac_objId => 12_0 => String
     * a => cc => 4 => String
     * b => tag => 5 => Long
     * c => name => 11 => String
     * d => cc1 => 15 => String
     * e => tag1 => 16 => Long
     * f => cc12 => 29 => String
     * g => tag12 => 30 => Long
     * h => cc1213 => 44 => String
     * i => tag1213 => 45 => Long
     * j => shapeLength => 66 => Double
     * k => shapeArea => 67 => Double
     * l => areaM2 => 68 => Double
     * @param jsc
     * @param z
     * @return
     */
    public JavaPairRDD<TileID, GeoItem> read(JavaSparkContext jsc, int z) {

        JavaRDD<GeoItem> geometry = jsc.textFile(datasource).map(new Function<String, GeoItem>() {
            @Override
            public GeoItem call(String s) throws Exception {
                String[] fields = s.split(Seprator.TAB.getValue());
                WKTReader wktReader = new WKTReader();
                Geometry geometry = wktReader.read(fields[fields.length-1]);
                String id = fields[12] + "_" + fields[0];
//                Map<String, Object> prop = new HashMap<>();
//                prop.put("a", fields[4]);
//                prop.put("b", fields[5].trim().length()>0 ? Long.valueOf(fields[5].trim()): null);
//                prop.put("c", fields[11]);
//                prop.put("d", fields[15]);
//                prop.put("e", fields[16].trim().length()>0 ? Long.valueOf(fields[16].trim()): null);
//                prop.put("f", fields[29]);
//                prop.put("g", fields[30].trim().length()>0 ? Long.valueOf(fields[30].trim()): null);
//                prop.put("h", fields[44]);
//                prop.put("i", fields[45].trim().length()>0 ? Long.valueOf(fields[45].trim()): null);
//                prop.put("j", fields[66].trim().length()>0 ? Double.valueOf(fields[66].trim()): null);
//                prop.put("k", fields[67].trim().length()>0 ? Double.valueOf(fields[67].trim()): null);
//                prop.put("l", fields[68].trim().length()>0 ? Double.valueOf(fields[68].trim()): null);
                GeoItem geoItem = new GeoItem();
                geoItem.setFid(id);
//                geoItem.setProp(prop);
                geoItem.setGeometry(JTS.transform(geometry,stTransform));
                return geoItem;
            }
        });

        // 每一个图斑单独处理，将其映射到瓦片中
        // key 为 tile ID，value 为对应的Geometry
        JavaPairRDD<TileID, GeoItem> tileGeometry = geometry.flatMapToPair(new PairFlatMapFunction<GeoItem, TileID, GeoItem>() {
            @Override
            public Iterator<Tuple2<TileID, GeoItem>> call(GeoItem geoItem) throws Exception {
                List<Tuple2<TileID, GeoItem>> result = new ArrayList<>();
                Geometry geometry = geoItem.getGeometry();
                // 获取geometry对应的envelope
                ReferencedEnvelope envelope = JTS.toEnvelope(geometry);
                // 获取geometry envelope对应的Tile范围
                ZLevelInfo tZLevelInfo = TileUtil.initZLevelInfoPZ(targetPyramidConfig, envelope)[z-mArgs.getZMin()];
                // 针对每一层级的不同瓦片，构造图斑处理的pipeline
                int tx_min = tZLevelInfo.getTileRanges()[0];
                int tx_max = tZLevelInfo.getTileRanges()[1];
                int ty_min = tZLevelInfo.getTileRanges()[2];
                int ty_max = tZLevelInfo.getTileRanges()[3];
                for (int tile_x = tx_min; tile_x <= tx_max; tile_x ++) {
                    for (int tile_y = ty_min; tile_y <= ty_max; tile_y ++) {
                        TileID tileID = new TileID();
                        tileID.setX(tile_x);
                        tileID.setY(tile_y);
                        tileID.setzLevel(z);
                        Envelope tileEnvelope = TileUtil.createTileBox(tileID, targetPyramidConfig);
                        Pipeline pipeline = getPipeline(CRS.parseWKT(mArgs.getTargetCRS()), BUFFER, tileEnvelope, false, true, true);
                        // 得到在每一个瓦片中对应的geometry
                        Geometry finalGeom = pipeline.execute(geometry);
                        if (!finalGeom.isEmpty()) {
                            Geometry pFinalGeom = onlyPolygon(finalGeom);
                            if (pFinalGeom == null) continue;
                            GeoItem finalItem = new GeoItem();
                            finalItem.setGeometry(pFinalGeom);
                            // geoItem.getProp().put("k", finalGeom.getArea());
                            // finalItem.setProp(geoItem.getProp());
                            finalItem.setFid(geoItem.getFid());
                            result.add(new Tuple2<>(tileID, finalItem));
                        }
                    }
                }
                return result.iterator();
            }
        });
        return tileGeometry;
    }

    public JavaPairRDD<TileID, GeoItem> buildTile(JavaPairRDD<TileID, GeoItem> tileGeometry) {
        // 根据瓦片ID group 对应的 geometry，并写出
        JavaPairRDD<TileID, GeoItem> out = tileGeometry.groupByKey().flatMapToPair(new TileBuilder());
        return out;
    }

    /**
     * 图斑重复有两种情况:
     * (1) 原始数据图斑重复
     * (2) 抽稀过程中，如果geometry完全在格网像素内，就会直接简化为格网像素的 envelope，造成很多完全重叠的小图斑
     * (3) 大的图斑在抽稀后,完全包含了小的图斑
     * 因此，需要去除result中，空间上完全重叠的图斑，基于 shape_area 字段, 保留面积最大的那一个
     */
    private List<Tuple2<TileID, GeoItem>> distinct(List<Tuple2<TileID, GeoItem>> in) {
        Map<String, Tuple2<TileID, GeoItem>> m = new HashMap<>();
        return new ArrayList<>(m.values());
    }

    private String getUniqKey(Tuple2<TileID, GeoItem> t) {
        return t._1.toString() + t._2.getFid();
    }

    private Geometry onlyPolygon(Geometry result) {
        if ((result instanceof Polygon) || (result instanceof MultiPolygon)) {
            return result;
        }
        List polys = org.locationtech.jts.geom.util.PolygonExtracter.getPolygons(result);
        if (polys.size() == 0) {
            return null;
        }
        if (polys.size() == 1) {
            return (Polygon) polys.get(0);
        }
        return new MultiPolygon(
                (Polygon[]) polys.toArray(new Polygon[polys.size()]), result.getFactory());
    }

    /**
     * 将裁切并抽稀的图斑合并
     * 如果直接合并,会有两个图斑不存在相交点的报错
     * @param g1
     * @param g2
     * @param buffer
     * @return
     */
    public Geometry union(Geometry g1, Geometry g2, double buffer, int depth, int maxDepth) {

        Geometry bg1;
        Geometry bg2;

        if (depth == 0) {
            bg1 = g1;
            bg2 = g2;
        } else {
            bg1 = g1.buffer(buffer);
            bg2 = g2.buffer(buffer);
        }

        try {
            return bg1.union(bg2);
        } catch (TopologyException e) {
            depth ++;
            if (depth > maxDepth) {
                return null;
            }
            return union(g1, g2, buffer*2, depth, maxDepth);
        }
    }

    private boolean isSmall(Tuple2<TileID, GeoItem> in) {
        Geometry g = in._2.getGeometry();
        Envelope te = TileUtil.createTileBox(in._1, targetPyramidConfig);
        double spanx = te.getWidth()/SCREEN_TILE_SIZE/2.0d;
        double spany = te.getHeight()/SCREEN_TILE_SIZE/2.0d;
        Envelope ge = g.getEnvelopeInternal();
        if (ge.getWidth() < spanx && ge.getHeight() < spany) {
            return true;
        }
        if (g.getArea() < (spanx*spany)) {
            return true;
        }
        return false;
    }

    public boolean equals(Geometry g1, Geometry g2) {
        int g1Num = g1.getNumGeometries();
        int g2Num = g2.getNumGeometries();
        if (g1Num != g2Num) {
            return false;
        }
        for (int n=0; n<g1Num; n++) {
            Geometry g1i = g1.getGeometryN(n);
            Geometry g2i = g2.getGeometryN(n);
            Coordinate[] g1c = g1i.getCoordinates();
            Coordinate[] g2c = g2i.getCoordinates();
            if (g1c.length != g2c.length)
                return false;
            for (int i=0; i<g1c.length; i++) {
                if (!g1c[i].equals(g2c[i]))
                    return false;
            }
        }
        return true;
    }

    /**
     * 构建瓦片中每个像素处理的Pipeline
     * 默认目标crs与源crs一致
     */
    private Pipeline getPipeline(CoordinateReferenceSystem sourceCrs, int buffer, Envelope referenceEnvelop, boolean preprocess, boolean isClip, boolean isSimplify) throws Exception {
        CoordinateReferenceSystem targetCrs = sourceCrs;
        int mapWidth = SCREEN_TILE_SIZE;
        int mapHeight = SCREEN_TILE_SIZE;
        int overSampleFactor = OVERSAMPLE_FACTOR;
        // rectangle of the image: width and height
        Rectangle paintArea = new Rectangle(mapWidth, mapHeight);

        // bounding box - in final map (target) CRS
        ReferencedEnvelope renderingArea = new ReferencedEnvelope(referenceEnvelop.getMinX(), referenceEnvelop.getMaxX(), referenceEnvelop.getMinY(), referenceEnvelop.getMaxY(), targetCrs);
        Pipeline pipeline = null;
        try {
            final PipelineBuilder builder =
                    PipelineBuilder.newBuilder(
                            renderingArea, paintArea, sourceCrs, overSampleFactor, buffer);

            if (preprocess){
                pipeline =
                        builder.preprocess().collapseCollections().precision(targetCrs).build();
            } else  {
                pipeline = builder.transform(false)
                        .clip(isClip, false)
                        .simplify(isSimplify)
                        .collapseCollections()
                        .precision(targetCrs)
                        .build();
            }
        } catch (FactoryException e) {
            e.printStackTrace();
        }

        return pipeline;
    }

    public static void main(String[] args) throws Exception {

        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setAppName("vector tile clipper");
        // conf.setMaster("yarn-cluster");
        conf.setMaster("local[*]");
        // conf.set("spark.es.resource", "index/type");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        IteratorVectorTileClipper clipper = new IteratorVectorTileClipper();
        clipper.init(args);

        int zMax = clipper.getMArgs().getZMax();
        logger.info("Initial the last level: " + zMax);
        JavaPairRDD<TileID, GeoItem> tileFeatures = clipper.read(jsc, zMax);
        tileFeatures = clipper.buildTile(tileFeatures);

        for (int i=zMax-1; i>=clipper.getMArgs().getZMin(); i--) {
            tileFeatures = clipper.buildTile(tileFeatures);
        }

        logger.info("Result: " + tileFeatures.count());
        logger.info("Model Finish");

    }

}
